{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "39d19711",
   "metadata": {},
   "source": [
    "# CW-Husky UserIO with JTAG\n",
    "\n",
    "One nice feature of the ChipWhisperer-Husky is the 20 pin UserIO connector, which is the front panel connector (not the side panel one). The connection of the UserIO data pins to equivalent JTAG pins is shown below for reference:\n",
    "\n",
    "| Userio Pin | JTAG Pin |\n",
    "|------------|----------|\n",
    "| UserIO[2]  | nRST     |\n",
    "| UserIO[3]  | TDO      |\n",
    "| UserIO[4]  | RTCK     |\n",
    "| UserIO[5]  | TCK      |\n",
    "| UserIO[6]  | TMS      |\n",
    "| UserIO[7]  | TDI      |\n",
    "\n",
    "For the following examples, you'll need to plug a second 20-pin cable from the UserIO connector to the JTAG connector. If you wanted you can *only* connect the front-side 20-pin cable as it also carries power, but normally you'd have the second one for regular communications.\n",
    "\n",
    "You'll also need to install [PyJtagBS](https://github.com/colinoflynn/pyjtagbs). The easiest way to do this is using pip with a pointer to the git repo:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "404df078",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install git+https://github.com/colinoflynn/pyjtagbs.git"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0276f7e0",
   "metadata": {},
   "source": [
    "With that running, you can now connect to your Husky as normal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03b760d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPE=\"OPENADC\"\n",
    "PLATFORM=\"CWHUSKY\"\n",
    "\n",
    "import chipwhisperer as cw\n",
    "scope = cw.scope(name='Husky')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43042c94",
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"../../Setup_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c6d5c24",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip uninstall pyjtagbs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67622770",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jtagbs import JTAGBS, JTAGCWUserIO\n",
    "\n",
    "interface = JTAGCWUserIO(scope)\n",
    "jtag = JTAGBS(interface)\n",
    "\n",
    "#jtag.init_scanchain()\n",
    "interface.scan_init_chain(True)\n",
    "\n",
    "print(jtag.list_devids())\n",
    "print(jtag.list_devices())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcfa0aa0",
   "metadata": {},
   "source": [
    "If the above worked, you should see this output:\n",
    "```\n",
    "IR Length: 3\n",
    "Found 1 devices\n",
    "[1268778103]\n",
    "[{'manid': '01000111011', 'pid': '1011101000000000', 'vn': '0100', 'manufacturer': 'ARM Ltd.'}]\n",
    "```\n",
    "\n",
    "Nice - you just connected over JTAG! This was just reading the device ID, from here there are a few things we can do.\n",
    "\n",
    "We'll now want to talk to the Arm Coresight register to check we can see that. To do this, you'll need to see the [Arm Coresight](https://developer.arm.com/documentation/ddi0314/h/Debug-Access-Port/JTAG-DP/Implementation-specific-details) documentation.\n",
    "\n",
    "From that documentation, you can see the instruction for reading the ID code is `b1110`, which is a 4-bit instruction width, and the associated data register read will be 32 bits. So the following should do those tasks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04ea0ee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "interface.write_IR(0b1110, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53c9993d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dr_value = interface.read_DR(32)\n",
    "print(dr_value)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "306bb836",
   "metadata": {},
   "source": [
    "This should return `[119, 4, 160, 75]`. This is LSB first, so you could construct a 32-bit value if you wanted with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b7cbd12",
   "metadata": {},
   "outputs": [],
   "source": [
    "import struct\n",
    "struct.unpack(\"<I\", bytes(dr_value))[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59ee3166",
   "metadata": {},
   "source": [
    "This should print the value of `1268778103`, which matches the ID printed on connection. Alright, let's keep going!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dd9cdc1",
   "metadata": {},
   "source": [
    "## Triggering on JTAG Commands\n",
    "\n",
    "One obviously interesting aspects of this will be if you can trigger on a JTAG command. Jtag commands are normally executed by the rising edge of the TCK pin. We can conviently trigger on the status of the userIO pins, so no hardware connections are needed! Set your trigger to D[5] (which is the TCK pin), making sure it's set to rising edge:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fce30669",
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.trigger.triggers = 'userio_d5'\n",
    "scope.adc.basic_mode = \"rising_edge\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "491156f3",
   "metadata": {},
   "source": [
    "You can use the argument `callback_before_updateir` to `write_IR`, which will call a given function before going into the `updateIR` state. You can use a similar function if you are writing the data register. If you want to trigger at other states, you'll just need to modify the source of `jtagbs` or implement your own function.\n",
    "\n",
    "Let's capture a power signature when we issue the command to do the ID code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0aabd5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "interface.write_IR(0b1110, 4, callback_before_updateir=scope.arm())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69df35ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.capture()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "195e35bf",
   "metadata": {},
   "source": [
    "You can see the length of the trigger: note that since this is bit-banged, it may be a very long value (like >500000 cycles!). It depends a bit on your computer, there is no specific good value here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12ec1662",
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.adc.trig_count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "022c090e",
   "metadata": {},
   "outputs": [],
   "source": [
    "trace = scope.get_last_trace()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8c7c105",
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.plot(trace)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19403ac7",
   "metadata": {},
   "source": [
    "This plot might not be too interesting, but if you were dealing with a JTAG unlock command it could be *very* interesting. Watch the clock in use: many devices won't use the external clock during JTAG, so you might want to turn off `hs2` (since it will just be adding tons of noise) & use a much faster ADC sampling rate.\n",
    "\n",
    "But you should experiment with your device to understand what clock source is being used. Some devices will use an external clock if present, which will give you a much more reliable power trace.\n",
    "\n",
    "Here is an example of turning off the clock & using a faster asynchronous sampling rate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7e2dc27",
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.io.hs2 = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36c03b56",
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.clock.clkgen_freq = 50E6\n",
    "scope.clock.adc_mul = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "406c0233",
   "metadata": {},
   "source": [
    "## Bonus - Low-Level Raw Bit-Banging\n",
    "\n",
    "How would you bing-bang a new protocol? Here's an example of some basic JTAG tests we did intially, which was just to put the JTAG port into bypass mode.\n",
    "\n",
    "To start with we set bit 5/6/7 (TCK, TMS, TDI) as outputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6812dbcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.userio.direction = 0b11100000"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c3ebd99",
   "metadata": {},
   "source": [
    "We then make some helper functions: one reads the state of the input pin, one writes data and toggles the clock:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab15ce82",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_tdo_status():\n",
    "    pins = scope.userio.status\n",
    "    if pins & (1<<3):\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "    \n",
    "def write(tms, tdi):\n",
    "    old = scope.userio.drive_data\n",
    "    old &= ~(1<<6 | 1<<7)\n",
    "    if tms:\n",
    "        old |= 1<<6\n",
    "    if tdi:\n",
    "        old |= 1<<7\n",
    "    \n",
    "    scope.userio.drive_data = old\n",
    "    scope.userio.drive_data = old | (1<<5)\n",
    "    scope.userio.drive_data = old & ~(1<<5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07988f68",
   "metadata": {},
   "source": [
    "The following is a basic test which forced devices into bypass mode, and tries to shift a single bit through the JTAG shift register. If the device is in bypass mode we should see that come out a clock cycle later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e79a42a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Having TMS high for 5 clock cycles is a JTAG TAP reset mode:\n",
    "write(1, 1)\n",
    "write(1, 1)\n",
    "write(1, 1)\n",
    "write(1, 1)\n",
    "write(1, 1)\n",
    "\n",
    "# This should get into the instruction register shift state\n",
    "write(0, 1) #\n",
    "write(1, 1)\n",
    "write(1, 1)\n",
    "write(0, 1)\n",
    "write(0, 1)\n",
    "\n",
    "# Send a bunch of 1's to force bypass mode - we don't know\n",
    "# how long the IR is potentially, so we could send many\n",
    "# more than 10. In this specific case I know it's not longer\n",
    "# than 10. We just need to keep TMS low to stay in the\n",
    "# SHIFT-IR state\n",
    "for i in range(0, 10):\n",
    "    write(0, 1)\n",
    "    \n",
    "# Exit shift-IR state so the instruction is used\n",
    "write(1,1)\n",
    "\n",
    "# Now go to shift-DR state\n",
    "write(1, 1)\n",
    "write(1, 1)\n",
    "write(0, 1)\n",
    "write(0, 1)\n",
    "\n",
    "# Write a bunch of 0's, the idea is to clear out the DR\n",
    "for i in range(0, 10):\n",
    "    write(0, 0)\n",
    "\n",
    "\n",
    "# Now read the output pin, sending a single 1 bit in.\n",
    "# If we are in bypass mode we should see a single 1 bit\n",
    "# come out, followed by a bunch of 0's\n",
    "for i in range(0, 10):\n",
    "    print(read_tdo_status())\n",
    "    if i == 0:\n",
    "        write(0, 1)\n",
    "    else:\n",
    "        write(0, 0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
